import os
from tqdm import tqdm
import xml.etree.ElementTree as ET
from sacremoses import MosesTokenizer
from pathlib import Path
import argparse


def moses_cut(in_file, out_file, lang):
    # 使用 MosesTokenizer 类来进行文本分词操作。
    # MosesTokenizer 是Moses机器翻译工具包中的一个组件，用于对文本进行分词处理。
    # 这里的 lang 参数用于指定使用的语言，如中文为 zh，英文为 en，日文为 ja。
    mt = MosesTokenizer(lang=lang)
    out_f = open(out_file, "w", encoding="utf8")
    with open(in_file, "r", encoding="utf8") as f:
        # 每读取一行，进行分词，并写入一行到新的文件中
        for line in f.readlines():
            line = line.strip()  # 去除字符串首尾的空白字符（包括空格、制表符和换行符等）
            if not line:  # 如果读入的行为空，则跳过
                continue
            # 对每行文本进行分词
            cut_line = mt.tokenize(line, return_str=True)  # 对每行文本进行分词
            out_f.write(cut_line.lower() + "\n")  # 变为小写，并写入文件
    out_f.close()


if __name__ == '__main__':
    # 创建了一个ArgumentParser对象，用于处理命令行参数的解析
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-p",
        "--pair_dir",
        default=None,
        type=str,
        help="The directory which contains language pair files.",
    )
    parser.add_argument(
        "-d",
        "--dest_dir",
        default=None,
        type=str,
        help="The destination directory to save processed train, dev and test file.",
    )
    parser.add_argument(
        "--src_lang",
        default="de",
        type=str,
        help="source language"
    )
    parser.add_argument(
        "--trg_lang",
        default="en",
        type=str,
        help="target language"
    )

    # 解析参数，args是一个列表，包含了传递的参数值
    args = parser.parse_args()
    # 如果不传参，就抛异常
    if not args.pair_dir:
        raise ValueError("Please specify --pair_dir")
    # 判断args.dest_dir是否存在,不存在就创建
    if not os.path.exists(args.dest_dir):
        os.makedirs(args.dest_dir)
    # 获取本地数据路径
    local_data_path = Path(args.pair_dir)
    # 获取保存路径
    data_dir = Path(args.dest_dir)

    # 分词
    for mode in ["train", "val", "test"]:
        moses_cut(
            local_data_path / f"{mode}.{args.src_lang}",  # 读取源语言文件
            data_dir / f"{mode}_src.cut.txt",
            lang=args.src_lang,
        )
        print(f"[{mode}] 源语言文本分词完成")
        moses_cut(
            local_data_path / f"{mode}.{args.trg_lang}",  # 读取目标语言文件
            data_dir / f"{mode}_trg.cut.txt",
            lang=args.trg_lang,
        )
        print(f"[{mode}] 目标语言文本分词完成")
